/*
 * Copyright (c) 2023-2023 elsfs Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.elsfs.cloud.spring.authorizationserver.component;

import com.nimbusds.jose.jwk.JWK;
import java.util.List;
import lombok.RequiredArgsConstructor;
import org.elsfs.cloud.api.security.key.OAuth2JWKRepository;
import org.elsfs.cloud.common.security.core.userdetails.SecurityUser;
import org.elsfs.cloud.spring.authorizationserver.federation.FederatedIdentityIdTokenCustomizer;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.core.oidc.StandardClaimNames;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.stereotype.Component;

/**
 * JSON Web密钥（JWK）源
 *
 * @author zeng
 */
@Component
@RequiredArgsConstructor
public class RsaKeyPairRepositoryOAuth2TokenCustomizer extends FederatedIdentityIdTokenCustomizer
    implements OAuth2TokenCustomizer<JwtEncodingContext> {
  private final OAuth2JWKRepository oauth2JWKRepository;

  @Override
  public void customize(JwtEncodingContext context) {
    List<JWK> keyPairs = this.oauth2JWKRepository.findJWKs();
    String kid = keyPairs.get(0).getKeyID();
    Authentication authentication = context.getPrincipal();
    if (authentication
        instanceof UsernamePasswordAuthenticationToken usernamePasswordAuthenticationToken) {
      Object principal = usernamePasswordAuthenticationToken.getPrincipal();
      if (principal instanceof SecurityUser securityUser) {
        addClaims(context.getClaims(), securityUser);
      }
    }
    context.getJwsHeader().keyId(kid);
    super.customize(context);
  }

  private void addClaims(JwtClaimsSet.Builder claims, SecurityUser securityUser) {
    claims.claim("username", securityUser.getUsername());
    claims.claim(StandardClaimNames.NICKNAME, securityUser.getNickname());
    claims.claim(StandardClaimNames.EMAIL, securityUser.getEmail());
  }
}
